package com.mm.cloth.rmbg;

import android.content.Context;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import java.io.IOException;
import java.io.InputStream;

/**
 * 基于Watermark-Removal ONNX的去水印处理器
 * 专门用于去除图像中的水印
 */
public class WatermarkRemovalProcessor {
    private static final String TAG = "WatermarkRemovalProcessor";
    
    private Context context;
    private WatermarkRemovalInferenceHelper inferenceHelper;
    private boolean isInitialized = false;
    
    // 模型文件路径
    private static final String MODEL_PATH = "watermark-removal.onnx";
    
    // 处理时间统计
    private float inferenceTime = 0;
    private float preprocessTime = 0;
    private float postprocessTime = 0;
    
    public WatermarkRemovalProcessor(Context context) {
        this.context = context;
        initializeProcessor();
    }
    
    /**
     * 初始化处理器
     */
    private void initializeProcessor() {
        try {
            Log.d(TAG, "初始化Watermark-Removal ONNX处理器");
            
            // 检查内存状态
            checkMemoryStatus();
            
            // 检查模型文件是否存在
            if (!checkModelFile()) {
                Log.e(TAG, "模型文件检查失败");
                isInitialized = false;
                return;
            }
            
            // 初始化推理助手
            inferenceHelper = new WatermarkRemovalInferenceHelper();
            
            // 加载模型
            loadWatermarkRemovalModel();
            
            isInitialized = true;
            Log.d(TAG, "Watermark-Removal处理器初始化成功");
            
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal处理器初始化失败", e);
            isInitialized = false;
            
            // 清理资源
            if (inferenceHelper != null) {
                try {
                    inferenceHelper.release();
                } catch (Exception ignored) {}
                inferenceHelper = null;
            }
        }
    }
    
    /**
     * 检查内存状态
     */
    private void checkMemoryStatus() {
        Runtime runtime = Runtime.getRuntime();
        long maxMemory = runtime.maxMemory();
        long totalMemory = runtime.totalMemory();
        long freeMemory = runtime.freeMemory();
        long usedMemory = totalMemory - freeMemory;
        long availableMemory = maxMemory - usedMemory;
        
        Log.d(TAG, "=== 内存状态检查 ===");
        Log.d(TAG, "最大内存: " + (maxMemory / (1024*1024)) + " MB");
        Log.d(TAG, "已用内存: " + (usedMemory / (1024*1024)) + " MB");
        Log.d(TAG, "可用内存: " + (availableMemory / (1024*1024)) + " MB");
        
        if (availableMemory < 100 * 1024 * 1024) {
            Log.w(TAG, "⚠️ 可用内存不足100MB，模型加载可能有风险");
        }
    }
    
    /**
     * 检查模型文件是否存在
     */
    private boolean checkModelFile() {
        try {
            InputStream modelStream = context.getAssets().open(MODEL_PATH);
            int size = modelStream.available();
            modelStream.close();
            
            Log.d(TAG, "模型文件检查通过: " + MODEL_PATH + ", 大小: " + (size / (1024*1024)) + " MB");
            return size > 0;
        } catch (Exception e) {
            Log.e(TAG, "模型文件检查失败: " + MODEL_PATH, e);
            return false;
        }
    }
    
    /**
     * 加载Watermark-Removal模型
     */
    private void loadWatermarkRemovalModel() {
        try {
            Log.d(TAG, "开始加载Watermark-Removal ONNX模型");
            
            AssetManager assetManager = context.getAssets();
            
            // 检查系统兼容性
            checkSystemCompatibility();
            
            // 尝试初始化ONNX Runtime模型
            inferenceHelper.initModel(assetManager, MODEL_PATH);
            
            // 验证初始化结果
            if (!inferenceHelper.isInitialized()) {
                throw new Exception("模型初始化验证失败");
            }
            
            Log.d(TAG, "Watermark-Removal ONNX模型加载成功");
            
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal模型加载失败", e);
            
            // 根据错误类型提供具体的错误信息
            String errorMessage;
            if (e.getMessage() != null && e.getMessage().contains("ONNX Runtime不可用")) {
                errorMessage = "ONNX Runtime库不可用";
            } else if (e.getMessage() != null && e.getMessage().contains("模型文件")) {
                errorMessage = "模型文件问题: " + e.getMessage();
            } else if (e.getMessage() != null && e.getMessage().contains("系统不兼容")) {
                errorMessage = "系统架构不兼容: " + e.getMessage();
            } else {
                errorMessage = "模型加载失败: " + e.getMessage();
            }
            
            throw new RuntimeException(errorMessage, e);
        }
    }
    
    /**
     * 检查系统兼容性
     */
    private void checkSystemCompatibility() throws Exception {
        String arch = System.getProperty("os.arch");
        String abi = android.os.Build.SUPPORTED_ABIS[0];
        
        Log.d(TAG, "系统架构: " + arch + ", ABI: " + abi);
        
        // 检查是否支持的架构
        if (!abi.contains("arm64") && !abi.contains("x86_64")) {
            throw new Exception("系统不兼容: 不支持的架构 " + abi);
        }
        
        // 检查Android版本
        int sdkVersion = android.os.Build.VERSION.SDK_INT;
        if (sdkVersion < 21) {
            throw new Exception("系统不兼容: Android版本过低 " + sdkVersion);
        }
        
        Log.d(TAG, "系统兼容性检查通过");
    }
    
    /**
     * 执行去水印处理
     * @param originalImage 原始图像
     * @param maskImage 水印遮罩图像（白色区域为水印位置）
     * @return 去水印后的图像
     */
    public Bitmap removeWatermark(Bitmap originalImage, Bitmap maskImage) {
        if (!isInitialized) {
            Log.e(TAG, "Watermark-Removal处理器未初始化");
            return null;
        }
        
        if (originalImage == null) {
            Log.e(TAG, "输入图像不能为空");
            return null;
        }
        
        Bitmap processedOriginal = null;
        Bitmap processedMask = null;
        Bitmap result = null;
        
        try {
            Log.d(TAG, "开始Watermark-Removal去水印处理");
            long totalStartTime = System.currentTimeMillis();
            
            // 检查内存状态
            System.gc(); // 建议垃圾回收
            
            // 预处理：确保图像格式正确
            long preprocessStart = System.currentTimeMillis();
            processedOriginal = preprocessImage(originalImage);
            
            // 如果有遮罩，处理遮罩；如果没有遮罩，使用自动检测
            if (maskImage != null) {
                processedMask = preprocessMask(maskImage, processedOriginal.getWidth(), processedOriginal.getHeight());
                Log.d(TAG, "使用用户提供的遮罩");
            } else {
                Log.d(TAG, "没有遮罩，使用自动水印检测");
            }
            
            long preprocessEnd = System.currentTimeMillis();
            preprocessTime = preprocessEnd - preprocessStart;
            
            Log.d(TAG, "预处理完成，开始Watermark-Removal推理");
            
            // 执行推理
            long inferenceStart = System.currentTimeMillis();
            if (processedMask != null) {
                // 使用遮罩的去水印
                result = inferenceHelper.removeWatermarkWithMask(processedOriginal, processedMask);
            } else {
                // 自动检测并去除水印
                result = inferenceHelper.removeWatermarkAuto(processedOriginal);
            }
            long inferenceEnd = System.currentTimeMillis();
            inferenceTime = inferenceEnd - inferenceStart;
            
            if (result == null) {
                Log.e(TAG, "Watermark-Removal推理返回空结果");
                return null;
            }
            
            Log.d(TAG, "Watermark-Removal推理完成，开始后处理");
            
            // 后处理
            long postprocessStart = System.currentTimeMillis();
            Bitmap finalResult = postprocessResult(result, originalImage);
            long postprocessEnd = System.currentTimeMillis();
            postprocessTime = postprocessEnd - postprocessStart;
            
            long totalEndTime = System.currentTimeMillis();
            Log.d(TAG, "Watermark-Removal去水印处理完成，总耗时: " + (totalEndTime - totalStartTime) + "ms");
            Log.d(TAG, "时间分解 - 预处理: " + preprocessTime + "ms, 推理: " + inferenceTime + "ms, 后处理: " + postprocessTime + "ms");
            
            return finalResult;
            
        } catch (OutOfMemoryError e) {
            Log.e(TAG, "Watermark-Removal处理内存不足", e);
            System.gc(); // 强制垃圾回收
            return null;
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal处理失败", e);
            return null;
        } finally {
            // 清理临时图像
            try {
                if (processedOriginal != null && processedOriginal != originalImage && !processedOriginal.isRecycled()) {
                    processedOriginal.recycle();
                }
                if (processedMask != null && processedMask != maskImage && !processedMask.isRecycled()) {
                    processedMask.recycle();
                }
            } catch (Exception cleanupException) {
                Log.w(TAG, "清理临时资源时出错", cleanupException);
            }
        }
    }
    
    /**
     * 预处理原始图像
     */
    private Bitmap preprocessImage(Bitmap originalImage) {
        // 确保图像格式为ARGB_8888
        if (originalImage.getConfig() != Bitmap.Config.ARGB_8888) {
            return originalImage.copy(Bitmap.Config.ARGB_8888, false);
        }
        return originalImage;
    }
    
    /**
     * 预处理遮罩图像
     */
    private Bitmap preprocessMask(Bitmap maskImage, int targetWidth, int targetHeight) {
        // 调整遮罩尺寸到目标尺寸
        Bitmap resizedMask = maskImage;
        if (maskImage.getWidth() != targetWidth || maskImage.getHeight() != targetHeight) {
            resizedMask = Bitmap.createScaledBitmap(maskImage, targetWidth, targetHeight, true);
        }
        
        // 确保遮罩格式为ARGB_8888
        if (resizedMask.getConfig() != Bitmap.Config.ARGB_8888) {
            Bitmap configMask = resizedMask.copy(Bitmap.Config.ARGB_8888, false);
            if (resizedMask != maskImage) {
                resizedMask.recycle();
            }
            resizedMask = configMask;
        }
        
        return resizedMask;
    }
    
    /**
     * 后处理结果
     */
    private Bitmap postprocessResult(Bitmap result, Bitmap originalImage) {
        // 如果结果尺寸与原图不同，调整到原图尺寸
        if (result.getWidth() != originalImage.getWidth() || result.getHeight() != originalImage.getHeight()) {
            return Bitmap.createScaledBitmap(result, originalImage.getWidth(), originalImage.getHeight(), true);
        }
        return result;
    }
    
    /**
     * 获取预处理时间
     */
    public float getPreprocessTime() {
        return preprocessTime;
    }
    
    /**
     * 获取推理时间
     */
    public float getInferenceTime() {
        return inferenceTime;
    }
    
    /**
     * 获取后处理时间
     */
    public float getPostprocessTime() {
        return postprocessTime;
    }
    
    /**
     * 释放资源
     */
    public void release() {
        try {
            if (inferenceHelper != null) {
                inferenceHelper.release();
                inferenceHelper = null;
            }
            
            // 强制垃圾回收
            System.gc();
            
            isInitialized = false;
            Log.d(TAG, "Watermark-Removal处理器资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放Watermark-Removal资源失败", e);
        }
    }
    
    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return isInitialized;
    }
}
